"""
Test script for Dung Beetle Optimizer (DBO) Algorithm
"""

import numpy as np
import matplotlib.pyplot as plt
from dbo import DungBeetleOptimizer


def run_comparison(function_name, function, dimensions=30, runs=10):
    """
    Run multiple tests on a function and report statistics
    
    Args:
        function_name: Name of the function
        function: The objective function
        dimensions: Number of dimensions
        runs: Number of independent runs
    """
    print(f"\n{'=' * 50}")
    print(f"Testing {function_name} function ({dimensions} dimensions)")
    print(f"{'=' * 50}")
    
    best_fitnesses = []
    
    # Set bounds based on function
    if function_name == "Rastrigin":
        lower_bound, upper_bound = -5.12, 5.12
    else:
        lower_bound, upper_bound = -30, 30
    
    # Run multiple times to get statistical results
    for run in range(runs):
        print(f"\nRun {run + 1}/{runs}:")
        
        optimizer = DungBeetleOptimizer(
            objective_function=function,
            dimensions=dimensions,
            population_size=50,
            max_iterations=100,
            lower_bound=lower_bound,
            upper_bound=upper_bound,
            rolling_factor=0.5,
            communication_probability=0.3
        )
        
        best_position, best_fitness, convergence = optimizer.optimize()
        best_fitnesses.append(best_fitness)
        
        print(f"Best fitness: {best_fitness}")
        
        # Only plot for the first run
        if run == 0:
            plt.figure(figsize=(10, 6))
            plt.plot(convergence)
            plt.title(f"Convergence Curve - {function_name} Function")
            plt.xlabel("Iteration")
            plt.ylabel("Fitness Value")
            plt.grid(True)
            plt.yscale("log")
            plt.savefig(f"{function_name}_convergence.png")
    
    # Report statistics
    best_fitnesses = np.array(best_fitnesses)
    print("\nStatistics:")
    print(f"Best fitness: {np.min(best_fitnesses)}")
    print(f"Worst fitness: {np.max(best_fitnesses)}")
    print(f"Mean fitness: {np.mean(best_fitnesses)}")
    print(f"Std dev: {np.std(best_fitnesses)}")


def sphere_function(x):
    """Sphere function (minimum at origin)"""
    return np.sum(x**2)


def rosenbrock_function(x):
    """Rosenbrock function"""
    return np.sum(100.0 * (x[1:] - x[:-1]**2)**2 + (x[:-1] - 1)**2)


def rastrigin_function(x):
    """Rastrigin function"""
    return 10 * len(x) + np.sum(x**2 - 10 * np.cos(2 * np.pi * x))


def griewank_function(x):
    """Griewank function"""
    sum_part = np.sum(x**2) / 4000
    prod_part = np.prod(np.cos(x / np.sqrt(np.arange(1, len(x) + 1))))
    return 1 + sum_part - prod_part


def ackley_function(x):
    """Ackley function"""
    term1 = -20 * np.exp(-0.2 * np.sqrt(np.sum(x**2) / len(x)))
    term2 = -np.exp(np.sum(np.cos(2 * np.pi * x)) / len(x))
    return term1 + term2 + 20 + np.exp(1)


if __name__ == "__main__":
    # Define test functions
    test_functions = {
        "Sphere": sphere_function,
        "Rosenbrock": rosenbrock_function,
        "Rastrigin": rastrigin_function,
        "Griewank": griewank_function,
        "Ackley": ackley_function
    }
    
    # Test dimensions
    dimensions = 30
    
    # Number of independent runs for statistical analysis
    runs = 5
    
    # Run tests for each function
    for name, func in test_functions.items():
        run_comparison(name, func, dimensions, runs)
